﻿#include<stdio.h>
#include<iostream>
#include<vector>
using namespace std;

enum COLOR {
	RED,
	BLACK,
};

template<class T>
struct RBTreeNode
{
	RBTreeNode(const T& data = T())
		:_data(data)
		, _pParent(nullptr)
		,_pLeft(nullptr)
		,_pRight(nullptr)
		, _color(RED)
	{}

	T _data;
	COLOR _color;
	RBTreeNode<T>* _pLeft;
	RBTreeNode<T>* _pRight;
	RBTreeNode<T>* _pParent;
};

template<class T>
class RBTree
{
	typedef RBTreeNode<T> Node;
public:
	RBTree()
	{
		_pHead = new Node;
		_pHead->_pLeft = _pHead;
		_pHead->_pRight = _pHead;
	}

	// 在红黑树中插入值为data的节点，插入成功返回true，否则返回false
	// 注意：为了简单起见，本次实现红黑树不存储重复性元素
	bool Insert(const T& data)
	{
		if (_pHead->_pParent == nullptr)
		{
			_pHead->_pParent = new Node(data);
			_pHead->_pParent->_color = BLACK;
			_pHead->_pLeft = _pHead->_pRight = _pHead->_pParent;
			_pHead->_pParent->_pParent=_pHead;
			return true;
		}
		
		//插入
		Node* child = nullptr;
		Node* cur = _pHead->_pParent;
		while (cur != nullptr)
		{
			if (data < cur->_data)
			{
				if (cur->_pLeft != nullptr)
					cur = cur->_pLeft;
				else
				{
					cur->_pLeft = new Node(data);
					child = cur->_pLeft;
					child->_pParent = cur;
					break;
				}
			}
			else if (data > cur->_data)
			{
				if (cur->_pRight != nullptr)
					cur = cur->_pRight;
				else
				{
					cur->_pRight = new Node(data);
					child = cur->_pRight;
					child->_pParent = cur;
					break;
				}
			}
			else
				return false;
		}

		Node* parent = child->_pParent;
		Node* grandpa = parent->_pParent;
		//变色或旋转
		while (parent->_color != BLACK)
		{
			if (grandpa->_pRight != nullptr && grandpa->_pLeft != nullptr)//有舅舅
			{
				if (grandpa->_pLeft->_color == grandpa->_pRight->_color)//舅舅红色
				{

					grandpa->_pLeft->_color = grandpa->_pRight->_color = BLACK;
					if (grandpa->_pParent != _pHead)
					{
						grandpa->_color = RED;
						child = grandpa;
						parent = child->_pParent;
						grandpa = parent->_pParent;
					}
					else
						break;


				}
				else//舅舅黑色
				{

					if (parent == grandpa->_pLeft)
					{
						RotateR(grandpa);
						parent->_color = BLACK;
						grandpa->_color = RED;
						break;
					}
					else
					{
						RotateL(grandpa);
						parent->_color = BLACK;
						grandpa->_color = RED;
						break;
					}

				}
			}
			else//没舅舅
			{

				if (parent == grandpa->_pLeft)
				{
					if (child == parent->_pLeft)
					{
						RotateR(grandpa);
						parent->_color = BLACK;
						grandpa->_color = RED;
					}

					else
					{
						RotateL(parent);
						RotateR(grandpa);
						child->_color = BLACK;
						grandpa->_color=RED;
					}
				}
				else
				{
					if (child == parent->_pRight)
					{
						RotateL(grandpa);
						parent->_color = BLACK;
						grandpa->_color = RED;
					}
					else
					{
						RotateR(parent);
						RotateL(grandpa);
						child->_color = BLACK;
						grandpa->_color = RED;
					}
				}
				break;
			}
		}


		//调整头结点指针
		if (_pHead->_pLeft->_pLeft != nullptr)//能运行到这至少插入了一个节点，树的最小值不会因为旋转改变
			_pHead->_pLeft = _pHead->_pLeft->_pLeft;

		if (_pHead->_pRight->_pRight != nullptr)
			_pHead->_pRight = _pHead->_pRight->_pRight;


		return true;
	}

	// 检测红黑树中是否存在值为data的节点，存在返回该节点的地址，否则返回nullptr
	Node* Find(const T& data)
	{
		Node* temp = _pHead->_pParent;
		while (temp != nullptr)
		{
			if (temp->_data == data)
				return temp;
			else if (data > temp->_data)
				temp = temp->_pRight;
			else
				temp = temp->_pLeft;
		}
		return nullptr;
	}

	
	// 获取红黑树最左侧节点
	Node* LeftMost()
	{
		return _pHead == nullptr ? nullptr : _pHead->_pLeft;
	}

	// 获取红黑树最右侧节点
	Node* RightMost()
	{
		return _pHead == nullptr ? nullptr : _pHead->_pRight;
	}

	// 检测红黑树是否为有效的红黑树，注意：其内部主要依靠_IsValidRBTRee函数检测
	bool IsValidRBTRee()
	{
		if (_pHead->_pParent == _pHead)
			return true;

		if (_pHead->_pParent->_color == RED)
			return false;

		int key = 0;
		Node* count = _pHead->_pParent;
		while (count != nullptr)
		{
			if (count->_color == BLACK)
				key++;
			count = count->_pLeft;
		}
		return _IsValidRBTRee(_pHead->_pParent, key);
	}

	void inorder(Node* root)
	{
		if (root == nullptr)
		{
			return;
		}

		if (root->_pLeft != nullptr)
			inorder(root->_pLeft);

		cout << root->_data<<"  ";
		if (root->_color == RED)
			cout << "RED" << "  ";
		else
			cout << "BLACK" << "  ";


		if (root->_pRight != nullptr)
			inorder(root->_pRight);
	}

	void InOrder()
	{
		Node* root = _pHead->_pParent;
		if (root == _pHead)
			return;
		inorder(root);
		cout << endl;
	}

private:


	bool _IsValidRBTRee(Node* pRoot, size_t blackCount, size_t pathBlack = 0)
	{
		if(pRoot == nullptr)
		{
			if (pathBlack == blackCount)
				return true;
			else
				return false;
		}
		if (pRoot->_color == BLACK)
			pathBlack++;

		bool a = true;
		if (pRoot->_pLeft != nullptr)
			a= _IsValidRBTRee(pRoot->_pLeft, blackCount, pathBlack);
		if (pRoot->_pRight != nullptr)
			a=  a && _IsValidRBTRee(pRoot->_pRight, blackCount, pathBlack);


		if (pRoot->_color == RED && pRoot->_pParent->_color == RED)
			return false;
		else
			return a;
	}




	// 左单旋
	void RotateL(Node* pParent)
	{
		Node* pChild = pParent->_pRight;
		pParent->_pRight = pChild->_pLeft;
		if(pParent->_pRight!=nullptr)
			pParent->_pRight->_pParent = pParent;
		pChild->_pLeft = pParent;
		pChild->_pParent = pParent->_pParent;
		if (pParent == pParent->_pParent->_pLeft)
		{
			pParent->_pParent->_pLeft = pChild;
		}
		else if (pParent == pParent->_pParent->_pRight)
		{
			pParent->_pParent->_pRight = pChild;
		}
		else//头结点
		{
			_pHead->_pParent = pChild;
		}

		pParent->_pParent = pChild;
	}

	// 右单旋
	void RotateR(Node* pParent)
	{
		Node* pChild = pParent->_pLeft;
		pParent->_pLeft = pChild->_pRight;
		if (pParent->_pRight != nullptr)
			pParent->_pLeft->_pParent = pParent;
		pChild->_pRight = pParent;
		pChild->_pParent = pParent->_pParent;

		if (pParent == pParent->_pParent->_pLeft)
		{
			pParent->_pParent->_pLeft = pChild;
		}
		else if (pParent == pParent->_pParent->_pRight)
		{
			pParent->_pParent->_pRight = pChild;
		}
		else//头结点
		{
			_pHead->_pParent = pChild;
		}

		pParent->_pParent = pChild;
	}
	// 为了操作树简单起见：获取根节点
	Node*& GetRoot()
	{
		return _pHead;

	}


private:
	Node* _pHead;
};


void TestAVLTree1()
{
	RBTree<int> t;
	// 常规的测试⽤例
	int a[] = { 16, 3, 7, 11, 9, 26, 18, 14, 15 };
	// 特殊的带有双旋场景的测试⽤例
	//int a[] = { 4, 2, 6, 1, 3, 5, 15, 0, 16, 14 };
	for (auto e : a)
	{
		t.Insert({ e });
	}
	t.InOrder();
	cout << t.IsValidRBTRee() << endl;
}

int main()
{
	TestAVLTree1();
	return 0;
}